"""
@Project    : cosmo-face
@Module     : train_centerface.py
@Author     : HuangJiWen[huangjiwen@haier.com]
@Created    : 2020/7/24 15:19
@Desc       :
"""

from __future__ import print_function

import argparse
import datetime
import os
import random
import time

import math
import torch
import torch.backends.cudnn as cudnn
import torch.nn.functional as F
import torch.optim as optim
import torch.utils.data as data

from data import WiderFaceDetection, detection_collate, preproc
from data import cfg_center_resnet18, cfg_center_resnet50, cfg_center_ghost, cfg_center_rfb
from layers.modules import centerpose_loss
from models.center_face.ghostnet_face import CenterGhostNetFace
from models.center_face.resnet_face import CenterResNetFace
from models.center_face.rfbnet_face import CenterRFBNetFace


def train(args, cfg):

    rgb_mean = (104, 117, 123)  # bgr order
    num_classes = 1
    img_dim = cfg['image_size']
    num_gpu = cfg['ngpu']
    batch_size = cfg['batch_size']
    max_epoch = cfg['epoch']
    gpu_train = cfg['gpu_train']

    net = None
    if args.network == "resnet18" or args.network == "resnet50":
        net = CenterResNetFace(cfg=cfg)
    elif args.network == "ghost":
        net = CenterGhostNetFace(cfg=cfg)
    elif args.network == "RFB":
        net = CenterRFBNetFace(cfg=cfg)
    else:
        print("Don't support network!")
        exit(0)

    print("Printing net...")
    print(net)

    if args.resume_net is not None:
        print('Loading resume network...')
        state_dict = torch.load(args.resume_net)
        # create new OrderedDict that does not contain `module.`
        from collections import OrderedDict
        new_state_dict = OrderedDict()

        if args.network == "RFB" or args.network == "resnet50":
            model_dict = net.state_dict()
            # 过滤操作
            # new_dict = {("backbone_model." + k): v for k, v in state_dict.items() if
            #             ("backbone_model." + k) in model_dict.keys()}
            new_dict = {k: v for k, v in state_dict.items() if k in model_dict.keys()}
            model_dict.update(new_dict)
            # 打印出来，更新了多少的参数
            print('Total : {}, update: {}'.format(len(state_dict), len(new_dict)))
            net.load_state_dict(model_dict)
            print("loaded finished!")
        elif args.network == "ghost":
            model_dict = net.state_dict()
            # 过滤操作
            new_dict = {("backbone_model." + k.replace("body.features1", "features").replace(
                "body.features2", "features").replace("body.features3", "features")): v
                        for k, v in state_dict.items() if
                        ("backbone_model." + k.replace("body.features1", "features").replace(
                            "body.features2", "features").replace("body.features3", "features")) in model_dict.keys()}
            model_dict.update(new_dict)
            # 打印出来，更新了多少的参数
            print('Total : {}, update: {}'.format(len(state_dict), len(new_dict)))
            net.load_state_dict(model_dict)
            print("loaded finished!")
        else:
            for k, v in state_dict.items():
                head = k[:7]
                if head == 'module.':
                    name = k[7:]  # remove `module.`
                else:
                    name = k
                new_state_dict[name] = v
            net.load_state_dict(new_state_dict)

    if torch.cuda.is_available():
        if num_gpu > 1 and gpu_train:
            net = torch.nn.DataParallel(net).cuda()
        else:
            net = net.cuda()

        cudnn.benchmark = True

    optimizer = optim.SGD(net.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay)
    criterion = centerpose_loss.CenterFaceLoss(num_classes)

    net.train()
    epoch = 0 + args.resume_epoch
    print('Loading Dataset...')

    dataset = WiderFaceDetection(args.training_dataset, preproc(img_dim, rgb_mean), img_dim=img_dim)

    epoch_size = math.ceil(len(dataset) / batch_size)
    max_iter = max_epoch * epoch_size

    stepvalues = (cfg['decay1'] * epoch_size, cfg['decay2'] * epoch_size)
    step_index = 0

    if args.resume_epoch > 0:
        start_iter = args.resume_epoch * epoch_size
    else:
        start_iter = 0

    for iteration in range(start_iter, max_iter):
        if iteration % epoch_size == 0:
            # create batch iterator
            batch_iterator = iter(data.DataLoader(dataset, batch_size, shuffle=True, num_workers=args.num_workers,
                                                  collate_fn=detection_collate))
            if (epoch % 10 == 0 and epoch > 0) or (epoch % 5 == 0 and epoch > cfg['decay1']):
                torch.save(net.state_dict(), args.save_folder + cfg['name'] + '_epoch_' + str(epoch) + '.pth')
            epoch += 1

        load_t0 = time.time()
        if iteration in stepvalues:
            step_index += 1
        lr = adjust_learning_rate(optimizer, args.lr, args.gamma, epoch, step_index, iteration, epoch_size)

        # load train data
        images, targets = next(batch_iterator)
        # images = images.float()

        # Multi-scale
        if args.multi_scale:
            sz = random.randrange(img_dim * 0.8, img_dim * 1.2 + 32) // 32 * 32  # size
            sf = sz / max(images.shape[2:])  # scale factor
            if sf != 1:
                ns = [math.ceil(x * sf / 32) * 32 for x in images.shape[2:]]  # new shape (stretched to gs-multiple)
                images = F.interpolate(images, size=ns, mode='bilinear', align_corners=False)

        if torch.cuda.is_available():
            images = images.cuda()
            targets = [anno.cuda() for anno in targets]

        # forward
        out = net(images)

        # back prop
        optimizer.zero_grad()
        loss_l, loss_c, loss_landm = criterion(out, targets, (out[0].shape[2], out[0].shape[3]))
        loss = loss_l + loss_c + loss_landm
        loss.backward()
        # torch.nn.utils.clip_grad_norm(net.parameters(), 256, norm_type=2)
        optimizer.step()
        load_t1 = time.time()
        batch_time = load_t1 - load_t0
        eta = int(batch_time * (max_iter - iteration))
        try:
            print('Epoch:{}/{} || Epochiter: {}/{} || Iter: {}/{} || Loc: {:.4f} Cla: {:.4f} Landm: {:.4f} || LR: {:.8f} || Batchtime: {:.4f} s || ETA: {}' \
                  .format(epoch, max_epoch, (iteration % epoch_size) + 1, epoch_size, iteration + 1, max_iter,
                          loss_l.item(), loss_c.item(), loss_landm.item(), lr, batch_time,
                          str(datetime.timedelta(seconds=eta))))
        except:
            print(loss_l, loss_c, loss_landm)

    torch.save(net.state_dict(), args.save_folder + cfg['name'] + '_Final.pth')


def adjust_learning_rate(optimizer, initial_lr, gamma, epoch, step_index, iteration, epoch_size):
    """Sets the learning rate
    # Adapted from PyTorch Imagenet example:
    # https://github.com/pytorch/examples/blob/master/imagenet/main.py
    """
    warmup_epoch = -1
    if epoch <= warmup_epoch:
        lr = 1e-6 + (initial_lr-1e-6) * iteration / (epoch_size * warmup_epoch)
    else:
        lr = initial_lr * (gamma ** step_index)
    for param_group in optimizer.param_groups:
        param_group['lr'] = lr
    return lr


if __name__ == '__main__':

    parser = argparse.ArgumentParser(description='Centerface Training')
    parser.add_argument('--training_dataset',
                        default='F:/gitee_project/cosmo-face/detection/data/wider_face/train/label.txt',
                        help='Training dataset directory')
    parser.add_argument('--network', default='resnet50', help='Backbone network resnet18 or resnet50 or ghost or RFB')
    parser.add_argument('--num_workers', default=0, type=int, help='Number of workers used in dataloading')
    parser.add_argument('--lr', '--learning-rate', default=2.5e-4, type=float, help='initial learning rate')
    parser.add_argument('--momentum', default=0.9, type=float, help='momentum')
    parser.add_argument('--resume_net',
                        # default=None,
                        # default="F:/gitee_project/cosmo-face/detection/weights/RFB_mosaic_epoch_245.pth",
                        # default="F:/gitee_project/cosmo-face/detection/weights/ghost_retina_epoch_210.pth",
                        default="F:/gitee_project/cosmo-face/detection/weights/resnet50_center_Final.pth",
                        help='resume net for retraining')
    parser.add_argument('--resume_epoch', default=0, type=int, help='resume iter for retraining')
    parser.add_argument('--weight_decay', default=5e-4, type=float, help='Weight decay for SGD')
    parser.add_argument('--gamma', default=0.1, type=float, help='Gamma update for SGD')
    parser.add_argument('--save_folder', default='./weights/', help='Location to save checkpoint models')
    parser.add_argument('--multi_scale', default=True, help='Location to save checkpoint models')

    args = parser.parse_args()

    if not os.path.exists(args.save_folder):
        os.mkdir(args.save_folder)

    cfg = None
    if args.network == "resnet18":
        cfg = cfg_center_resnet18
    elif args.network == "resnet50":
        cfg = cfg_center_resnet50
    elif args.network == "ghost":
        cfg = cfg_center_ghost
    elif args.network == "RFB":
        cfg = cfg_center_rfb
    else:
        print("Don't support network!")
        exit(0)

    train(args, cfg)
